import torch.nn

import shared.torchtools


class PinnOfKortewegDeVriesEquation(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.network: torch.nn.Sequential = torch.nn.Sequential(
            torch.nn.Linear(2, 300),
            torch.nn.Tanh(),
            *[
                 torch.nn.Linear(300, 300),
                 torch.nn.Tanh()
             ] * 3,
            torch.nn.Linear(300, 1)
        )
        self.lambda1: torch.nn.Parameter = torch.nn.Parameter(
            torch.tensor(0.), requires_grad=True)
        self.lambda2: torch.nn.Parameter = torch.nn.Parameter(
            torch.tensor(0.), requires_grad=True)

    def forward(self, t: torch.Tensor, x: torch.Tensor):
        tx = torch.stack([t, x], dim=-1)
        f = self.network(tx)
        return f[:, 0]

    def physics_loss(self, point_count: int = 1000) -> torch.Tensor:
        t = shared.torchtools.rand(
            point_count, 0, 1, True)
        x = shared.torchtools.rand(
            point_count, -1, 1, True)

        u = self(t, x)

        du_dt, du_dx = shared.torchtools.grad(u, [t, x], True, True)
        du_dx_dx, = shared.torchtools.grad(du_dx, [x], True, True)
        du_dx_dx_dx, = shared.torchtools.grad(du_dx_dx, [x], True, True)

        loss = du_dt + self.lambda1 * u * du_dx + self.lambda2 * du_dx_dx_dx
        return torch.mean(loss ** 2)
